import networkx as nx
import numpy as np

class connect:
  def __init__(self):
    self.nb = []
    self.w  = []
    self.num = 0
  def oneStep(self,pos):
    r = np.zeros(pos.shape,dtype=np.complex128)
    for i in range(len(pos)):
      r[i] = np.average(pos[self.nb[i]],axis=0,weights=self.w[i])
    return r
  def dump(self,N):
    for l in range(N):
      print("n",self.nb[l],"w",self.w[l])
  def remove(self,k,n):
    for i in range(len(self.nb[k])):
      if self.nb[k][i] == n:
        self.w[k][i] = 0
class egress(connect):
  def __init__(self,G):
    connect.__init__(self)
    for n in range(nx.number_of_nodes(G)):
      r0 = [n]
      w0 = [1]
      for nb in nx.all_neighbors(G,n):
        weight = G.get_edge_data(n,nb)
        if weight!=None:
          ww = weight['weight']
          w0.append(ww)
          r0.append(nb)
      self.nb.append(r0)
      if len(w0)>1:
        w0[0] = 0.1
      self.w.append(w0)
      self.num += 1

class two(connect):
  def __init__(self,G):
    connect.__init__(self)
    for n in range(nx.number_of_nodes(G)):
      nbs = {}
      for nb in nx.all_neighbors(G,n):
        weight = G.get_edge_data(n,nb)
        if weight!=None:
          for nbb in nx.all_neighbors(G,nb):
            weight2 = G.get_edge_data(nb,n)
            if weight2!=None:
              if nbs.__contains__(nbb):
                nbs[nbb] += weight2['weight']*weight['weight']
              else:
                nbs[nbb] = weight2['weight']*weight['weight']
      r0 = []
      w0 = []
      for k in nbs:
        r0.append(k)
        w0.append(nbs[k])
      if len(r0)==0:
        r0 = [n]
        w0 = [1]
      index = np.array(w0).argsort()
      if len(index)>20:
        index = index[:20] 
      r0 = np.array(r0)[index].tolist()
      w0 = np.array(w0)[index].tolist()
      self.nb.append(r0)
      self.w.append(w0)
      self.num += 1
  
class ingress(connect):
  def __init__(self,G):
    connect.__init__(self)
    for n in range(nx.number_of_nodes(G)):
      r0 = [n]
      w0 = [1]
      for nb in nx.all_neighbors(G,n):
        weight = G.get_edge_data(nb,n)
        if weight!=None:
          ww = weight['weight']
          w0.append(ww)
          r0.append(nb)
      self.nb.append(r0)
      if len(w0)>1:
        w0[0] = 0.1
      self.w.append(w0)
      self.num += 1
      